/*-
 * Copyright 2003-2005 Colin Percival
 * All rights reserved
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted providing that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
 * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
 * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
 * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING
 * IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */
#include <bzlib.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>

//#include <err.h>
//#include <unistd.h>
//#include <fcntl.h>
#include <bzlib.h>


//#define _RTT
#define TRUE  (1)
#define FALSE (0)

#ifdef _RTT
#include <rtthread.h>
#include <dfs_posix.h>
#define memcpy rt_memcpy
#define malloc rt_malloc
#endif

#define err(level,...) \
        do{\
			printf(__VA_ARGS__);\
			exit(0);\
		}while(0)

#define errx(level,...) \
		do{\
			printf(__VA_ARGS__);\
			exit(0);\
		}while(0)

/// memory file operation API///////
typedef enum
{
    MF_OK,
    MF_ERROR,
    MF_EOF,
    MF_SEEK_SET,
    MF_SEEK_CUR,
    MF_SEEK_END,
} MF_STATE;

typedef struct memfile
{
    char* buf;
    int   len;
    int   pos;
    int   stat;
    int   oflag;
} MEM_FILE;

MEM_FILE* mf_open(void *mem,size_t len);
void mf_close    (MEM_FILE *mf);

int mf_read (MEM_FILE *mf,void *buf,size_t len);
int mf_write(MEM_FILE *mf,void *buf,size_t len);
int mf_error(MEM_FILE *mf);
int mf_eof  (MEM_FILE *mf);
int mf_putc (MEM_FILE *mf,int c);
int mf_getc (MEM_FILE *mf);
int mf_tell (MEM_FILE* mf);
void mf_rewind(MEM_FILE* mf);
int mf_seek (MEM_FILE* mf,off_t off,int whence);

MEM_FILE* mf_open(void *mem,size_t len)
{
    MEM_FILE* M = (MEM_FILE*)malloc(sizeof(MEM_FILE));
    if(M==NULL)
        return NULL;

    M->buf = (char*)mem;
    M->len = len;
    M->pos = 0;
    M->stat = MF_OK;
    return M;
}

void mf_close (MEM_FILE *mf)
{
    free(mf);
}

int mf_tell (MEM_FILE* mf)
{
    return mf->pos;
}

void mf_rewind(MEM_FILE* mf)
{
    mf->pos = 0;
    mf->stat = MF_OK;
}

int mf_seek (MEM_FILE* mf,off_t off,int whence)
{
    switch(whence)
    {
    case MF_SEEK_SET:
        mf->pos = off;
        if(mf->pos > mf->len)
        {
            mf->pos = mf->len;
        }
        break;

    case MF_SEEK_CUR:
        mf->pos += off;
        if(mf->pos > mf->len)
        {
            mf->pos = mf->len;
        }
        break;

    case MF_SEEK_END:
        mf->pos = mf->len;
        break;
    }
    return 0;
}

int mf_read (MEM_FILE *mf,void *buf,size_t len)
{
    int actual = 0;

    if(mf->pos + len <= mf->len)
    {
        memcpy(buf,&mf->buf[mf->pos],len);
        actual = len;
    }
    else if(mf->pos + len > mf->len)
    {
        actual = mf->len - mf->pos;
        memcpy(buf,&mf->buf[mf->pos],actual);
    }

    mf->pos += actual;

    if(mf->pos == mf->len)
        mf->stat = MF_EOF;

    return actual;
}

int mf_error(MEM_FILE *mf)
{
    return 0; // always
}

int mf_eof  (MEM_FILE *mf)
{
    return mf->stat == MF_EOF;
}

#define BZ_SETERR(eee)                    \
{                                         \
   if (bzerror != NULL) *bzerror = eee;   \
   if (bzf != NULL) bzf->lastErr = eee;   \
}

// in memory file object
typedef struct
{
    MEM_FILE*     handle;
    char          buf[BZ_MAX_UNUSED];
    ssize_t       bufN;
    int           writing;
    bz_stream     strm;
    size_t        lastErr;
    int           initialisedOk;
} bzMemFile;

/*---------------------------------------------------*/
static BZFILE* BZ2_bzMemReadOpen
( int*  bzerror,
  MEM_FILE* f,
  int   verbosity,
  int   small,
  void* unused,
  int   nUnused )
{
    bzMemFile* bzf = NULL;
    int     ret;

    BZ_SETERR(BZ_OK);

    if (f == NULL ||
            (small != 0 && small != 1) ||
            (verbosity < 0 || verbosity > 4) ||
            (unused == NULL && nUnused != 0) ||
            (unused != NULL && (nUnused < 0 || nUnused > BZ_MAX_UNUSED)))
    {
        BZ_SETERR(BZ_PARAM_ERROR);
        return NULL;
    };

    if (mf_error(f))
    {
        BZ_SETERR(BZ_IO_ERROR);
        return NULL;
    };

    bzf = malloc ( sizeof(bzMemFile) );
    if (bzf == NULL)
    {
        BZ_SETERR(BZ_MEM_ERROR);
        return NULL;
    };

    BZ_SETERR(BZ_OK);

    bzf->initialisedOk = FALSE;
    bzf->handle        = f;
    bzf->bufN          = 0;
    bzf->writing       = FALSE;
    bzf->strm.bzalloc  = NULL;
    bzf->strm.bzfree   = NULL;
    bzf->strm.opaque   = NULL;

    while (nUnused > 0)
    {
        bzf->buf[bzf->bufN] = *((char*)(unused));
        bzf->bufN++;
        unused = ((void*)( 1 + ((char*)(unused))  ));
        nUnused--;
    }

    ret = BZ2_bzDecompressInit ( &(bzf->strm), verbosity, small );
    if (ret != BZ_OK)
    {
        BZ_SETERR(ret);
        free(bzf);
        return NULL;
    };

    bzf->strm.avail_in = bzf->bufN;
    bzf->strm.next_in  = bzf->buf;

    bzf->initialisedOk = TRUE;
    return bzf;
}

/*---------------------------------------------------*/
static void BZ2_bzMemReadClose ( int *bzerror, BZFILE *b )
{
    bzMemFile* bzf = (bzMemFile*)b;

    BZ_SETERR(BZ_OK);
    if (bzf == NULL)
    {
        BZ_SETERR(BZ_OK);
        return;
    };

    if (bzf->writing)
    {
        BZ_SETERR(BZ_SEQUENCE_ERROR);
        return;
    };

    if (bzf->initialisedOk)
        (void)BZ2_bzDecompressEnd ( &(bzf->strm) );
    free ( bzf );
}

/*---------------------------------------------------*/
static int BZ2_bzMemRead
( int*    bzerror,
  BZFILE* b,
  void*   buf,
  int     len )
{
    size_t  ret;
    size_t  n;
    bzMemFile* bzf = (bzMemFile*)b;

    BZ_SETERR(BZ_OK);

    if (bzf == NULL || buf == NULL || len < 0)
    {
        BZ_SETERR(BZ_PARAM_ERROR);
        return 0;
    };

    if (bzf->writing)
    {
        BZ_SETERR(BZ_SEQUENCE_ERROR);
        printf("bz 01\n");
        return 0;
    };

    if (len == 0)
    {
        BZ_SETERR(BZ_OK);
        return 0;
    };

    bzf->strm.avail_out = len;
    bzf->strm.next_out = buf;

    while (1)
    {

        if (mf_error(bzf->handle))
        {
            BZ_SETERR(BZ_IO_ERROR);
            return 0;
        };

        if (bzf->strm.avail_in == 0 && !mf_eof(bzf->handle))
        {
            n = mf_read (bzf->handle, bzf->buf, sizeof(char)*BZ_MAX_UNUSED);
            if (mf_error(bzf->handle))
            {
                BZ_SETERR(BZ_IO_ERROR);
                printf("bz 02\n");
                return 0;
            };
            bzf->bufN = (size_t)n;
            bzf->strm.avail_in = bzf->bufN;
            bzf->strm.next_in = bzf->buf;
        }

        ret = BZ2_bzDecompress ( &(bzf->strm) );

        if (ret != BZ_OK && ret != BZ_STREAM_END)
        {
            BZ_SETERR(ret);
            printf("bz 03 ret=%d\n",ret);
            return 0;
        };

        if (ret == BZ_OK && mf_eof(bzf->handle) &&
                bzf->strm.avail_in == 0 && bzf->strm.avail_out > 0)
        {
            BZ_SETERR(BZ_UNEXPECTED_EOF);
            printf("bz 04\n");
            return 0;
        };

        if (ret == BZ_STREAM_END)
        {
            BZ_SETERR(BZ_STREAM_END);
            return len - bzf->strm.avail_out;
        };
        if (bzf->strm.avail_out == 0)
        {
            BZ_SETERR(BZ_OK);
            return len;
        };

    }
    return 0; /*not reached*/
}


static off_t offtin(u_char *buf)
{
    off_t y;

    y=buf[7]&0x7F;
    y=y*256;
    y+=buf[6];
    y=y*256;
    y+=buf[5];
    y=y*256;
    y+=buf[4];
    y=y*256;
    y+=buf[3];
    y=y*256;
    y+=buf[2];
    y=y*256;
    y+=buf[1];
    y=y*256;
    y+=buf[0];

    if(buf[7]&0x80) y=-y;

    return y;
}

int patch_mem(char* oldfile,ssize_t old_len,
              char* patchfile,ssize_t patch_len,
              char** newfile,ssize_t *new_len)
{
    MEM_FILE  *f,
              *cpf,
              *dpf,
              *epf;

    BZFILE * cpfbz2,
           * dpfbz2,
           * epfbz2;

    int cbz2err,
        dbz2err,
        ebz2err;

    ssize_t oldsize = old_len,newsize;
    ssize_t bzctrllen,bzdatalen;
    char header[32],buf[8];
    char *oldfile_mem = oldfile, *newfile_mem;
    off_t oldpos,newpos;
    off_t ctrl[3];
    off_t lenread;
    off_t i;

    /*
    File format:
    	0	8	"BSDIFF40"
    	8	8	X
    	16	8	Y
    	24	8	sizeof(newfile)
    	32	X	bzip2(control block)
    	32+X	Y	bzip2(diff block)
    	32+X+Y	???	bzip2(extra block)
    with control block a set of triples (x,y,z) meaning "add x bytes
    from oldfile to x bytes from the diff block; copy y bytes from the
    extra block; seek forwards in oldfile by z bytes".
    */
    f = mf_open(patchfile,patch_len);
    /* Read header */
    if (mf_read(f,header, 32) < 32)
    {
        if (mf_eof(f))
            errx(1, "00 Corrupt patch\n");
    }

    /* Check for appropriate magic */
    if (memcmp(header, "BSDIFF40", 8) != 0)
        errx(1, "01 Corrupt patch\n");

    /* Read lengths from header */
    bzctrllen = offtin(header + 8);
    bzdatalen = offtin(header + 16);
    newsize   = offtin(header + 24);
    if((bzctrllen<0) || (bzdatalen<0) || (newsize<0))
        errx(1,"02 Corrupt patch\n");

    printf("control data len=%d\n",bzctrllen);
    printf("diff data len=%d\n",bzdatalen);
    printf("newsize data len=%d\n",newsize);

    /* Close patch file and re-open it via libbzip2 at the right places */
    mf_close(f);

    /******* control data block *********/
    if ((cpf = mf_open(patchfile, patch_len)) == NULL)
        err(1, "mf_open control\n");

    if (mf_seek(cpf, 32, MF_SEEK_SET))            // control offset is 32
        err(1, "mf_seek cpf 32\n");

    if ((cpfbz2 = BZ2_bzMemReadOpen(&cbz2err, cpf, 0, 0, NULL, 0)) == NULL)
        errx(1, "BZ2_bzReadOpen, bz2err = %d", cbz2err);

    /******* diff data block *********/
    if ((dpf = mf_open(patchfile, patch_len)) == NULL)
        err(1, "mf_open diff");

    if (mf_seek(dpf, 32 + bzctrllen, MF_SEEK_SET))
        err(1, "mf_seek dpf %lld\n",(long long)(32 + bzctrllen));

    if ((dpfbz2 = BZ2_bzMemReadOpen(&dbz2err, dpf, 0, 0, NULL, 0)) == NULL)
        errx(1, "BZ2_bzReadOpen, bz2err = %d\n", dbz2err);

    /******* extra data block *********/
    if ((epf = mf_open(patchfile, patch_len)) == NULL)
        err(1, "mf_open extra\n");

    if (mf_seek(epf, 32 + bzctrllen + bzdatalen, MF_SEEK_SET))
        err(1, "mf_seek epf %lld\n",(long long)(32 + bzctrllen + bzdatalen));

    if ((epfbz2 = BZ2_bzMemReadOpen(&ebz2err, epf, 0, 0, NULL, 0)) == NULL)
        errx(1, "BZ2_bzReadOpen, bz2err = %d", ebz2err);

    // this is to be written into flash
    if((newfile_mem = malloc(newsize + 1)) == NULL)
        err(1,NULL);

    oldpos=0;
    newpos=0;
    while(newpos<newsize)
    {
        /* Read control data */
        for(i=0; i<=2; i++)
        {
            lenread = BZ2_bzMemRead(&cbz2err, cpfbz2, buf, 8);

            if ((lenread < 8) || ((cbz2err != BZ_OK) &&
                                  (cbz2err != BZ_STREAM_END)))
            {
                errx(1, "03 Corrupt patch  lenread:%d\n",lenread);
            }

            ctrl[i]=offtin(buf);
        };

        /* Sanity-check */
        if( newpos+ctrl[0] > newsize)
        {
            errx(1,"04 Corrupt patch\n");
        }

        /* Read diff string */
        lenread = BZ2_bzMemRead(&dbz2err, dpfbz2, newfile_mem + newpos, ctrl[0]);
        if ((lenread < ctrl[0]) ||
                ((dbz2err != BZ_OK) && (dbz2err != BZ_STREAM_END)))
            errx(1, "05 Corrupt patch\n");

        /* Add old data to diff string */
        for(i=0; i<ctrl[0]; i++)
        {
            if((oldpos + i >= 0) && (oldpos + i < oldsize))
                newfile_mem[newpos + i] += oldfile_mem[oldpos + i];
        }
        /* Adjust pointers */
        newpos += ctrl[0];
        oldpos += ctrl[0];

        /* Sanity-check */
        if( newpos + ctrl[1] > newsize)
        {
            errx(1,"06 Corrupt patch\n");
        }

        /* Read extra string */
        lenread = BZ2_bzMemRead(&ebz2err, epfbz2, newfile_mem + newpos, ctrl[1]);
        if ((lenread < ctrl[1]) ||
                ((ebz2err != BZ_OK) && (ebz2err != BZ_STREAM_END)))
        {
            errx(1, "07 Corrupt patch\n");
        }

        /* Adjust pointers */
        newpos += ctrl[1];
        oldpos += ctrl[2];
    };

    /* Clean up the bzip2 reads */
    BZ2_bzMemReadClose(&cbz2err, cpfbz2);
    BZ2_bzMemReadClose(&dbz2err, dpfbz2);
    BZ2_bzMemReadClose(&ebz2err, epfbz2);
    mf_close(cpf);
    mf_close(dpf);
    mf_close(epf);

    *newfile  = newfile_mem;
    *new_len  = newsize;
    return 0;
}

int main(int argc,char** argv)
{
    FILE *oldf,
         *newf,
         *patchf;

    char *buf_old,
         *buf_new,
         *buf_pat;

    long newlen;

    if(argc < 3)
    {
        printf("use: oldfile newfile patchfile\n");
        return 0;
    }

    oldf = fopen(argv[1],"rb");
    fseek(oldf,0,SEEK_END);
    long oldlen = ftell(oldf);
    buf_old = malloc(oldlen);
    printf("open old file:%s oldlen=%d\n",argv[0],oldlen);
    rewind(oldf);
    fread(buf_old,1,oldlen,oldf);


    patchf = fopen(argv[3],"rb");
    fseek(patchf,0,SEEK_END);
    long patlen = ftell(patchf);
    printf("open patch file:%s patlen=%d\n",argv[3],patlen);
    buf_pat = malloc(patlen);
    rewind(patchf);
    fread(buf_pat,1,patlen,patchf);


    newf = fopen(argv[2],"wb");

    patch_mem(buf_old,oldlen,
              buf_pat,patlen,
              &buf_new,&newlen);

    printf("final: newlen=%d\n",newlen);

    fwrite(buf_new,1,newlen,newf);

    fclose(newf);
    free(buf_old);
    free(buf_pat);
    free(buf_new);

    fclose(oldf);
    fclose(patchf);
    return 0;
}
